from flair.embeddings.document import TransformerDocumentEmbeddings
from flair.data import Sentence
from sentence_transformers import SentenceTransformer, models
import torch

text = "Good Morning"
document_embeddings = TransformerDocumentEmbeddings(model="bert-base-uncased", pooling="mean")

# option 1
sentence = Sentence(text, use_tokenizer=False)
document_embeddings.embed(sentence)
embedding_1 = sentence.get_embedding()
print(embedding_1.shape)
print(embedding_1[:20])

# option 2
document_embeddings.pooling = 'max'
sentence = Sentence(text, use_tokenizer=False)
document_embeddings.embed(sentence)
embedding_2 = sentence.get_embedding()
print(embedding_2.shape)
print(embedding_2[:20])

# option 3
document_embeddings.pooling = 'cls'
sentence = Sentence(text, use_tokenizer=False)
document_embeddings.embed(sentence)
embedding_3 = sentence.get_embedding()
print(embedding_3.shape)
print(embedding_3[:20])


word_embedding_model = models.Transformer('bert-base-uncased', max_seq_length=512)
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())

model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

# option 1'
embedding_11 = model.encode(text, convert_to_tensor=True)
print(embedding_11.shape)
print(embedding_11[:20])

# option 2'
pooling_model.pooling_mode_cls_token = False
pooling_model.pooling_mode_mean_tokens = False
pooling_model.pooling_mode_max_tokens = True
embedding_21 = model.encode(text, convert_to_tensor=True)
print(embedding_21.shape)
print(embedding_21[:20])

# option 3'
pooling_model.pooling_mode_cls_token = True
pooling_model.pooling_mode_mean_tokens = False
pooling_model.pooling_mode_max_tokens = False
embedding_31 = model.encode(text, convert_to_tensor=True)
print(embedding_31.shape)
print(embedding_31[:20])

assert torch.all(torch.eq(embedding_1, embedding_11))
assert torch.all(torch.eq(embedding_2, embedding_21))
assert torch.all(torch.eq(embedding_3, embedding_31))